package com.playwright.util.gpt;


import cn.hutool.core.util.ObjectUtil;
import cn.hutool.http.HttpUtil;
import cn.hutool.http.Method;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;

import java.net.InetSocketAddress;
import java.net.Proxy;
import java.util.*;
import java.util.stream.Collectors;


/**
 * ChatGPT 测试.
 *
 * @author zenghao
 * @date 2023/2/26
 */
public class ChatGPTTest {

    private static final String API_KEY = "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";

    private static final String API_DOMAIN = "https://api.openai.com";

    public static void main(String[] args) {
        System.setProperty("jdk.tls.useExtendedMasterSecret", "false");
        System.setProperty("jdk.tls.client.protocols", "TLSv1.2");
        System.setProperty("https.protocols", "TLSv1.2");

        dialog();
    }

    private static void dialog() {
        Scanner scanner = new Scanner(System.in);
        System.out.println("请开始输入你的问题：");
        String text = scanner.next();

        LinkedList<ChatCompletion> chats = new LinkedList<>();
        while (!Objects.equals("exit", text)) {
            chats.add(ChatCompletion.ofUser(text));

            System.out.println("你的输入：");
            System.out.println(text);

            String reply = chatCompletions(chats);
            System.out.println("ChatGPT回复：");
            System.out.println(reply);
            chats.add(ChatCompletion.ofAssistant(reply));

            if (chats.size() > 6) {
                chats.removeFirst();
            }

            System.out.println("-------------------------------");
            System.out.println("输入你的问题（或输入【exit】退出）：");
            text = scanner.next();
        }
    }

    public static String chatCompletions(List<ChatCompletion> chats) {
        Map<String, Object> param = new HashMap<>();
        param.put("model", "gpt-3.5-turbo");
        param.put("messages", chats);
        JSONObject jsonObject = execute("/v1/chat/completions", Method.POST, param);

        JSONArray jsonArray = jsonObject.getJSONArray("choices");
        if (ObjectUtil.isEmpty(jsonArray)) {
            return "";
        }
        List<ChatCompletion> messages = jsonArray.stream().map(JSONUtil::parseObj)
                .map(j -> JSONUtil.toBean(j.getJSONObject("message"), ChatCompletion.class))
                .collect(Collectors.toList());
        return messages.get(0).getContent();
    }

    public static JSONObject execute(String url, Method method, Map<String, Object> param) {
        String request = JSONUtil.toJsonStr(param);
        //System.out.println("request = " + request);
        Proxy proxy = new Proxy(Proxy.Type.HTTP, new InetSocketAddress("127.0.0.1", 7890));
        String response = HttpUtil.createRequest(method, API_DOMAIN + url).setProxy(proxy)
                .header("Authorization", "Bearer " + API_KEY)
                .body(request).execute().body();

        //System.out.println("response = " + response);
        return JSONUtil.parseObj(response);
    }
}
